import torch
import torch.nn as nn


class Deep_Linear_Model(nn.Module):
    """
    input_size = 4
    hidden_size = 8
    num_layers = 2
    output_size = 1
    device = 'cuda'
    """

    def __init__(self, output_size=1):
        super().__init__()
        self.nn_sequential = nn.Sequential(
            nn.Linear(5, 32),
            nn.Linear(32, 32),
            nn.ReLU(),
            nn.Linear(32, 64),
            # nn.BatchNorm1d(64),

            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            # nn.GELU(),
            nn.Linear(64, 128),
            nn.Linear(128, output_size)
        )

    def forward(self, x):
        out = self.nn_sequential(x)
        return out


def deepLinearModel(output_size):
    return Deep_Linear_Model(output_size)
